import sys
from typing import Dict, List
from common.model.milvus_utils import check_database_exists
from pymilvus import DataType, MilvusClient

class PaperSegmentModel():
    """论文切割文档向量数据库"""
    
    _DB_NAME = "paper_segment"
    _COLLECTION_NAME = "paper_segment_collection"
    _META_COLLECTION_NAME = "paper_meta_collection"
    
    INDEX_VECTOR_DIM = 1024
    TEXT_CHUNK_SIZE  = 2048
    
    def __init__(self, milvus_client: MilvusClient) -> None:
        self._milvus_client = milvus_client
        
        self._init_db()
        # todo: 临时存储方案，后续可以考虑将这块挪到paperfile里存
        self._init_collection(PaperSegmentModel._META_COLLECTION_NAME, self._create_meta_collection)
        self._init_collection(PaperSegmentModel._COLLECTION_NAME, self._create_collection)
        
    def _init_db(self) -> None:
        if not check_database_exists(self._milvus_client, PaperSegmentModel._DB_NAME):
            self._milvus_client.create_database(
                db_name=PaperSegmentModel._DB_NAME
            )
        
        self._milvus_client.use_database(
            db_name=PaperSegmentModel._DB_NAME
        )
        
    def _init_collection(self, collection_name: str, create_func) -> None:
        if not self._milvus_client.has_collection(collection_name):
            create_func()
        
        self._milvus_client.load_collection(
            collection_name=collection_name
        )

        res = self._milvus_client.get_load_state(
            collection_name=collection_name
        )

        print(f'load milvus db {PaperSegmentModel._DB_NAME} collection {collection_name}: {res}', file=sys.stderr)
    
    def _create_meta_collection(self) -> None:
        schema = MilvusClient.create_schema(
            auto_id=False,
            enable_dynamic_field=True,
        )

        schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
        schema.add_field(field_name="paper_basic_info", datatype=DataType.VARCHAR, max_length=4096)
        schema.add_field(field_name="abstract_summary", datatype=DataType.VARCHAR, max_length=4096)
        schema.add_field(field_name="user_id", datatype=DataType.INT64)
        schema.add_field(field_name="paper_file_id", datatype=DataType.VARCHAR, max_length=128)
        schema.add_field(field_name="nothing", datatype=DataType.FLOAT_VECTOR, dim=2)
        
        index_params = MilvusClient.prepare_index_params()

        index_params.add_index(
            field_name="user_id",
            index_name="idx_user_id",
            index_type="STL_SORT"
        )
        
        index_params.add_index(
            field_name="paper_file_id",
            index_name="idx_paper_file_id",
            index_type="INVERTED"
        )
        
        index_params.add_index(
            field_name="nothing",
            index_name="idx_nothing",
            index_type="FLAT",
            metric_type="COSINE",
            params={}
        )
        
        # 创建集合
        collection_name = PaperSegmentModel._META_COLLECTION_NAME
        self._milvus_client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)
        
    def insert_meta(self, user_id: int, paper_file_id: str, paper_basic_info: str, abstract_summary: str) -> Dict:
        """插入论文meta数据"""
        data = [
            {
                "user_id": user_id,
                "paper_file_id": paper_file_id,
                "paper_basic_info": paper_basic_info,
                "abstract_summary": abstract_summary,
                "nothing": [-1, -1],
            }
        ]
        
        res = self._milvus_client.insert(
            collection_name=self._META_COLLECTION_NAME, 
            data=data
        )
            
        return res
    
    def find_meta(self, user_id: int, paper_file_id: str) -> Dict | None:
        """查找论文meta数据"""
        filter_expression = f"user_id == {user_id} AND paper_file_id == '{paper_file_id}'"
        res = self._milvus_client.query(
            collection_name=self._META_COLLECTION_NAME,
            filter=filter_expression,
            output_fields=["*"]
        )
        
        if len(res) > 0:
            return res[0]
        else:
            return None
    
    def delete_meta(self, user_id: int, paper_file_id: str) -> int:
        """根据user_id和paper_file_id删除论文meta数据"""
        filter_expression = f"user_id == {user_id} AND paper_file_id == '{paper_file_id}'"
        res = self._milvus_client.delete(
            collection_name=self._META_COLLECTION_NAME,
            filter=filter_expression
        )
            
        if isinstance(res, list):
            deleted_count = len(res)  # 主键ID列表的长度即为删除数量
        elif hasattr(res, 'delete_count'):  # 检查是否为OmitZeroDict类型
            deleted_count = res.delete_count
        else:
            deleted_count = -1  # 默认处理未知类型
            
        return deleted_count
    
    def _create_collection(self) -> None:
        schema = MilvusClient.create_schema(
            auto_id=False,
            enable_dynamic_field=True,
        )

        schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
        schema.add_field(field_name="embedding_vector", datatype=DataType.FLOAT_VECTOR, dim=PaperSegmentModel.INDEX_VECTOR_DIM)
        schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=PaperSegmentModel.TEXT_CHUNK_SIZE)
        schema.add_field(field_name="meta", datatype=DataType.JSON)
        schema.add_field(field_name="splitter_id", datatype=DataType.VARCHAR, max_length=128)
        schema.add_field(field_name="user_id", datatype=DataType.INT64)
        schema.add_field(field_name="paper_file_id", datatype=DataType.VARCHAR, max_length=128)
        schema.add_field(field_name="paper_file_name", datatype=DataType.VARCHAR, max_length=128)
        
        index_params = MilvusClient.prepare_index_params()

        index_params.add_index(
            field_name="user_id",
            index_name="idx_user_id",
            index_type="STL_SORT"
            # index_type="INVERTED"
        )
        
        index_params.add_index(
            field_name="paper_file_id",
            index_name="idx_paper_file_id",
            index_type="INVERTED"
        )
            
        index_params.add_index(
            field_name="embedding_vector",
            index_name="idx_vector_flat",
            index_type="FLAT",
            metric_type="COSINE",
            params={}
        )
        
        # 创建集合
        collection_name = PaperSegmentModel._COLLECTION_NAME
        self._milvus_client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)
        
    def paper_segment_data(self, user_id: int, paper_file_id: str, paper_file_name: str, 
               text: str, embedding_vector: list, splitter_id: str, meta: Dict=None) -> Dict:
        """插入论文片段数据"""
        return {
            "user_id": user_id,
            "paper_file_id": paper_file_id,
            "paper_file_name": paper_file_name,
            "text": text,
            "embedding_vector": embedding_vector,
            "meta": meta if meta else {},
            "splitter_id": splitter_id
        }
        
    def insert_one(self, user_id: int, paper_file_id: str, paper_file_name: str, 
               text: str, embedding_vector: list, splitter_id: str, meta: Dict=None) -> Dict:
        """插入论文片段数据"""
        data = [
            self.paper_segment_data(user_id, paper_file_id, paper_file_name, 
               text, embedding_vector, splitter_id, meta)
        ]
        
        res = self._milvus_client.insert(
            collection_name=self._COLLECTION_NAME, 
            data=data
        )
            
        return res
    
    def insert_many(self, datas: List[Dict]) -> Dict:
        """插入论文片段数据"""
        return self._milvus_client.insert(
            collection_name=self._COLLECTION_NAME, 
            data=datas
        )

    def delete_by(self, user_id: int, paper_file_id: str) -> int:
        """根据user_id和paper_file_id组合条件删除数据"""
        filter_expression = f"user_id == {user_id} AND paper_file_id == '{paper_file_id}'"
        res = self._milvus_client.delete(
            collection_name=self._COLLECTION_NAME,
            filter=filter_expression
        )
            
        if isinstance(res, list):
            deleted_count = len(res)  # 主键ID列表的长度即为删除数量
        elif hasattr(res, 'delete_count'):  # 检查是否为OmitZeroDict类型
            deleted_count = res.delete_count
        else:
            deleted_count = -1  # 默认处理未知类型
            
        return deleted_count

    def search_by(self, user_id: int, paper_file_id: str, 
                    query_vectors: list, top_k: int=5) -> List[List[dict]]:
        """先按组合条件过滤，再向量检索"""
        # 第一步：组合条件过滤
        filter_expression = f"user_id == {user_id} AND paper_file_id == '{paper_file_id}'"
        filter_res = self._milvus_client.search(
            collection_name=self._COLLECTION_NAME,
            data=[query_vectors],
            anns_field="embedding_vector",
            search_params={"params": {}},  # No additional parameters required for FLAT
            limit=top_k,
            filter=filter_expression,
            output_fields=["text", "meta", "splitter_id"]
        )
            
        return filter_res

    def check_exist(self, user_id: int, paper_file_id: str) -> bool:
        """检查user_id和paper_file_id组合条件是否存在数据"""
        filter_expression = f"user_id == {user_id} AND paper_file_id == '{paper_file_id}' RANDOM_SAMPLE(0.01)"
        res = self._milvus_client.query(
            collection_name=self._COLLECTION_NAME,
            filter=filter_expression
        )
            
        return len(res) > 0